package udtf;

import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDTF;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.StructObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;

import java.util.ArrayList;
import java.util.List;

/**
 * UDTF substr_words(string str, int length)
 * example:
 * desc: substr_words('abcde', 2), return explode(array<string>), string={ab, bc, cd, de}
 */
public class SubStringUDTF extends GenericUDTF {

    private final List<String> outList = new ArrayList<>();

    @Override
    public StructObjectInspector initialize(StructObjectInspector argOIs) {
        List<String> fieldName = new ArrayList<>();
        List<ObjectInspector> fieldOIs = new ArrayList<>();

        fieldName.add("sentence");
        fieldOIs.add(PrimitiveObjectInspectorFactory.javaStringObjectInspector);

        return ObjectInspectorFactory.getStandardStructObjectInspector(fieldName, fieldOIs);
    }

    @Override
    public void process(Object[] args) throws HiveException {
        if (args == null || args.length != 2) {
             return;
        }
        String string = args[0].toString();
        int length = Integer.parseInt(args[1].toString());
        if (length <= 0 || string.length() < length) {
            return;
        }

        // ("abcdef", 2) ==> (ab, bc, cd, de, ef)
        for (int i = 0; i <= string.length()-length; i++) {
            outList.clear();
            outList.add(string.substring(i, i+length)); // beginLength(from 0, 1, 2, 3, 4), endLength(from 2, 3, 4, 5, 6)
            forward(outList);
        }
    }

    @Override
    public void close() throws HiveException {
    }
}
